# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
# with the following modifications:
# - It uses the patched version of `sde_step_with_logprob` from `sd3_sde_with_logprob.py`.
# - It returns all the intermediate latents of the denoising process as well as the log probs of each denoising step.
from typing import Any, Dict, List, Optional, Union
import torch
from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
from .sd3_sde_with_logprob import sde_step_with_logprob

@torch.no_grad()
def pipeline_with_logprob(
    self,  # 重写StableDiffusion3Pipeline类的方法
    prompt: Union[str, List[str]] = None,
    prompt_2: Optional[Union[str, List[str]]] = None,
    prompt_3: Optional[Union[str, List[str]]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 28,
    sigmas: Optional[List[float]] = None,
    guidance_scale: float = 7.0,
    negative_prompt: Optional[Union[str, List[str]]] = None,
    negative_prompt_2: Optional[Union[str, List[str]]] = None,
    negative_prompt_3: Optional[Union[str, List[str]]] = None,
    num_images_per_prompt: Optional[int] = 1,
    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_prompt_embeds: Optional[torch.FloatTensor] = None,
    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
    negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
    clip_skip: Optional[int] = None,
    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
    max_sequence_length: int = 256,
    skip_layer_guidance_scale: float = 2.8,
    noise_level: float = 0.7,
):
    height = height or self.default_sample_size * self.vae_scale_factor
    width = width or self.default_sample_size * self.vae_scale_factor

    # 1. 检查输入参数是否正确，如果不正确则抛出错误
    self.check_inputs(
        prompt, # 第一个提示词
        prompt_2, # 第二个提示词（用于多提示词扩展，非必需）
        prompt_3, # 第三个提示词（用于多提示词扩展，非必需）
        height, # 生成图像的高度
        width, # 生成图像的宽度
        negative_prompt=negative_prompt, # 第一个负面提示词
        negative_prompt_2=negative_prompt_2, # 第二个负面提示词
        negative_prompt_3=negative_prompt_3, # 第三个负面提示词
        prompt_embeds=prompt_embeds, # 已经编码好的提示词嵌入（可选）
        negative_prompt_embeds=negative_prompt_embeds, # 已编码的负面提示词嵌入（可选）
        pooled_prompt_embeds=pooled_prompt_embeds, # 全局池化后的提示词嵌入（可选）
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, # 全局池化后的负面提示词嵌入（可选）
        callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, # 每一步结束时回调的张量输入（可选）
        max_sequence_length=max_sequence_length, # 提示词最大序列长度（用于截断或填充）
    )


    self._guidance_scale = guidance_scale # 引导尺度，用于控制生成图像与提示词的相关性
    self._skip_layer_guidance_scale = skip_layer_guidance_scale # 跳过层引导尺度，用于控制跳过层的影响
    self._clip_skip = clip_skip # CLIP跳过层数，用于控制CLIP模型的处理
    self._joint_attention_kwargs = joint_attention_kwargs # 联合注意力的参数（可选）
    self._interrupt = False # 中断标志，用于在生成过程中中断操作

    # 2. 配置相关参数
    if prompt is not None and isinstance(prompt, str):
        batch_size = 1 # 如果提示词是字符串，则批次大小为1
    elif prompt is not None and isinstance(prompt, list):
        batch_size = len(prompt) # 如果提示词是列表，则批次大小为列表长度
    else:
        batch_size = prompt_embeds.shape[0] # 如果没有提供提示词，则从嵌入中获取批次大小

    device = self._execution_device
    
    # 获取LoRA缩放因子，如果联合注意力参数中有scale则使用它，否则为None
    lora_scale = (self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None) 
    
    # 3. 编码提示词
    (
        prompt_embeds,
        negative_prompt_embeds,
        pooled_prompt_embeds,
        negative_pooled_prompt_embeds,
    ) = self.encode_prompt(
        prompt=prompt,
        prompt_2=prompt_2,
        prompt_3=prompt_3,
        negative_prompt=negative_prompt,
        negative_prompt_2=negative_prompt_2,
        negative_prompt_3=negative_prompt_3,
        do_classifier_free_guidance=self.do_classifier_free_guidance,
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        device=device,
        clip_skip=self.clip_skip,
        num_images_per_prompt=num_images_per_prompt,
        max_sequence_length=max_sequence_length,
        lora_scale=lora_scale,
    )
    # 如果使用classifier_free_guidance, 则将负面提示词嵌入和正面提示词嵌入合并
    if self.do_classifier_free_guidance:
        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
        pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)

    # 4. 准备 latent variables | 使用 VAE 编码器将图像压缩成 latent 表示（潜空间）
    num_channels_latents = self.transformer.config.in_channels
    latents = self.prepare_latents(
        batch_size * num_images_per_prompt,
        num_channels_latents,
        height,
        width,
        prompt_embeds.dtype,
        device,
        generator,
        latents,
    )

    # 5. Prepare timesteps | 扩散模型的采样过程准备时间步信息
    scheduler_kwargs = {}
    timesteps, num_inference_steps = retrieve_timesteps(
        self.scheduler,
        num_inference_steps,
        device,
        sigmas=sigmas,
        **scheduler_kwargs,
    )
    num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
    self._num_timesteps = len(timesteps)

    # 6. Prepare image embeddings
    all_latents = [latents] # 用于存储每一步的latent
    all_log_probs = [] # 用于存储每一步的log概率

    # 7. Denoising loop | 迭代去噪过程
    with self.progress_bar(total=num_inference_steps) as progress_bar: #num_inference_steps 10
        
        # 迭代每个时间步
        for i, t in enumerate(timesteps): 
            '''
            tensor([1000.0000,  960.1293,  913.3490,  857.6923,  790.3683,  707.2785, 602.1506,  464.8760,  278.0488,    8.9286]
            '''
            if self.interrupt:
                continue
            
            # 1. 噪声预测
            latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents # 如果使用classifier free guidance，则扩展 latent 的维度为 2 倍
            timestep = t.expand(latent_model_input.shape[0]) # 扩展每一个时间步(共10个时间步)到 batch 维度 (batchsize*2)
            
            # import pdb
            # pdb.set_trace()
            noise_pred = self.transformer(  # 使用DiT模型进行噪声预测
                hidden_states=latent_model_input, # vae编码后的latent | torch.Size([8*2, 16, 64, 64])
                timestep=timestep, # 当前时间步 | torch.Size([8*2]) 
                encoder_hidden_states=prompt_embeds, # 编码后的prompt嵌入 | torch.Size([8*2, 205, 4096])
                pooled_projections=pooled_prompt_embeds, # 全局的prompt嵌入 | torch.Size([8*2, 2048])
                joint_attention_kwargs=self.joint_attention_kwargs, # 联合注意力参数
                return_dict=False,
            )[0]
            noise_pred = noise_pred.to(prompt_embeds.dtype) # torch.Size([8, 16, 64, 64])
           
           # classifier_free_guidance
            if self.do_classifier_free_guidance: 
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) # 则将 噪声预测 分为两部分：无条件(noise_pred_uncond)和有条件(noise_pred_text)
                noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) # 使用引导尺度调整噪声预测
                
            latents_dtype = latents.dtype

            # 2. 去噪 | 为了适用GRPO, 原来的 ODE 改为 SDE
            latents, log_prob, prev_latents_mean, std_dev_t = sde_step_with_logprob( # 使用SDE步骤进行去噪
                self.scheduler, # 扩散模型的调度器
                noise_pred.float(), # 预测的噪声 ｜ v theta
                t.unsqueeze(0), # 当前时间步
                latents.float(), # 当前的latent
                noise_level = noise_level, # 噪声水平
            )
            
            # 3. 记录每一步去噪得到的latent和log概率
            all_latents.append(latents)
            all_log_probs.append(log_prob)
            if latents.dtype != latents_dtype:
                latents = latents.to(latents_dtype)
            
            # call the callback, if provided
            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                progress_bar.update()

    latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
    latents = latents.to(dtype=self.vae.dtype)
    image = self.vae.decode(latents, return_dict=False)[0]
    image = self.image_processor.postprocess(image, output_type=output_type)

    # Offload all models
    self.maybe_free_model_hooks()

    return image, all_latents, all_log_probs
